package io.gitee.kenewstar.dbmq.data;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.date.DateTime;
import cn.hutool.core.date.DateUtil;
import cn.hutool.core.util.ArrayUtil;
import cn.hutool.core.util.RandomUtil;
import cn.hutool.db.Db;
import cn.hutool.db.Entity;
import io.gitee.kenewstar.dbmq.util.Const;
import io.gitee.kenewstar.dbmq.ConsumerRegisterInfo;
import io.gitee.kenewstar.dbmq.ConsumerTask;
import io.gitee.kenewstar.dbmq.DbMqException;
import io.gitee.kenewstar.dbmq.Message;
import io.gitee.kenewstar.dbmq.util.KeyWordMqField;
import io.gitee.kenewstar.dbmq.util.MySqlTableSql;
import io.gitee.kenewstar.dbmq.util.Util;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.sql.DataSource;
import java.net.InetAddress;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.List;
import java.util.Objects;

/**
 * @author kenewstar
 */
public class DbMySqlHandler implements DbDataHandler, DbDataHandlerAdapter {

    private static final Logger LOGGER = LoggerFactory.getLogger(DbMySqlHandler.class);

    private DataSource dataSource;

    public DbMySqlHandler() {}

    public DbMySqlHandler(DataSource dataSource) {
        this.dataSource = dataSource;
    }

    private Db db() {
        return Db.use(dataSource);
    }

    @Override
    public DbDataHandler adapter() {
        return this;
    }

    @Override
    public String getType() {
        return DbDataHandlerWrapper.MYSQL;
    }

    @Override
    public void initConsumerRegister() {
        final String consumerRegisterSql = MySqlTableSql.build().getConsumerRegisterSql();
        try {
            db().execute(consumerRegisterSql);
        } catch (Exception e) {
            throw new DbMqException(e);
        }
    }

    @Override
    public void initConsumerQueue(List<ConsumerTask> taskList) {
        final MySqlTableSql tableSql = MySqlTableSql.build();
        try {
            for (ConsumerTask task : taskList) {
                db().execute(tableSql.createQueueTableSql(task.getTopic()));
            }
        } catch (Exception e) {
            throw new DbMqException(e);
        }
    }

    @Override
    public List<ConsumerRegisterInfo> pullRegisterInfo(List<String> topics) {
        if (CollUtil.isEmpty(topics)) {
            return CollUtil.newArrayList();
        }
        final Entity where = Entity.create(MySqlTableSql.TABLE_CONSUMER_REGISTER);
        where.set(KeyWordMqField.TOPIC, topics);
        try {
            return db().findAll(where, ConsumerRegisterInfo.class);
        } catch (Exception e) {
            throw new DbMqException(e);
        }
    }

    @Override
    public ConsumerRegisterInfo createConsumer(ConsumerTask task) {
        try {
            Entity consumer = Entity.create(MySqlTableSql.TABLE_CONSUMER_REGISTER)
                    .set(KeyWordMqField.TOPIC, task.getTopic())
                    .set(KeyWordMqField.IP, Util.getIpAddress())
                    .set(KeyWordMqField.UUID, task.getUuid())
                    .set(KeyWordMqField.HEARTBEAT_TIME, DateUtil.date());
            task.setId(db().insertForGeneratedKey(consumer));
            return consumer.toBean(ConsumerRegisterInfo.class);
        } catch (Exception e) {
            throw new DbMqException(e);
        }
    }

    @Override
    public void updateConsumer(ConsumerTask task, List<ConsumerRegisterInfo> registerInfo) {
        final DateTime current = DateUtil.date();
        String updateSql = "update " + MySqlTableSql.TABLE_CONSUMER_REGISTER +
                " set heartbeat_time = ?, ip = ?, uuid = ?, version = version + 1, status = 'up' " +
                " where id = ? and uuid = ? and version = ?";
        int result = 0;
        try {
            for (ConsumerRegisterInfo info : registerInfo) {
                result = db().execute(updateSql, current, InetAddress.getLocalHost().getHostAddress(), task.getUuid(),
                        info.getId(), info.getUuid(), info.getVersion());
                // 替换成功
                if (result == 1) {
                    task.setOldId(task.getId());
                    task.setId(info.getId());
                    break;
                }
            }
            // 替换失败
            if (result == 0) {
                this.createConsumer(task);
            }
        } catch (Exception e) {
            throw new DbMqException(e);
        }
    }

    @Override
    public void registerDown(Long id) {
        if (Objects.isNull(id)) {
            return;
        }
        try {
            Entity where = Entity.create(MySqlTableSql.TABLE_CONSUMER_REGISTER).set(KeyWordMqField.ID, id);
            db().update(new Entity().set(KeyWordMqField.STATUS, Const.Status.DOWN), where);
        } catch (Exception e) {
            throw new DbMqException(e);
        }
    }

    @Override
    public void registerDown(List<ConsumerRegisterInfo> list) {
        if (CollUtil.isEmpty(list)) {
            return;
        }
        try {
            String downSql = "update `" + MySqlTableSql.TABLE_CONSUMER_REGISTER +
                    "` set status = 'down', version = version + 1" +
                    " where id = ? and status = 'up' and version = ? and heartbeat_time = ?";
            for (ConsumerRegisterInfo info : list) {
                db().execute(downSql, info.getId(), info.getVersion(), info.getHeartbeatTime());
            }
        } catch (Exception e) {
            throw new DbMqException(e);
        }
    }

    @Override
    public void sendHeartbeat(List<Long> ids) {
        if (CollUtil.isEmpty(ids)) {
            return;
        }

        LOGGER.info("=====================Send Heartbeat===========================");
        try {
            Entity where = Entity.create(MySqlTableSql.TABLE_CONSUMER_REGISTER)
                    .set(KeyWordMqField.ID, ArrayUtil.toArray(ids, Long.class));
            db().update(new Entity().set(KeyWordMqField.HEARTBEAT_TIME, DateUtil.date()).set(KeyWordMqField.STATUS, Const.Status.UP), where);
        } catch (Exception e) {
            throw new DbMqException(e);
        }
    }

    @Override
    public List<Message> getMessage(ConsumerTask task) {
        if (Objects.isNull(task.getTag())) {
            return CollUtil.newArrayList();
        }
        try {
            String sql = "select * from `" + task.getTopic() + "` where consume_flag = ? and consumer_tag = ? limit ?";
            return db().query(sql, Message.class, 0, task.getTag(), task.getLimit());
        } catch (Exception e) {
            throw new DbMqException(e);
        }
    }

    @Override
    public void deleteMessage(String topic, List<Long> ids) {
        if (CollUtil.isEmpty(ids)) {
            return;
        }

        try {
            db().tx(db -> db.del(topic, KeyWordMqField.ID, ArrayUtil.toArray(ids, Long.class)));
        } catch (Exception e) {
            throw new DbMqException(e);
        }
    }

    @Override
    public int sendMessage(String topic, List<String> messageBodyList, List<Integer> tags) {
        try {
            final DateTime current = DateUtil.date();
            List<Entity> list = new ArrayList<>(messageBodyList.size());
            for (String messageBody : messageBodyList) {
                Entity entity = Entity.create(topic);
                entity.set(KeyWordMqField.MESSAGE_BODY, messageBody);
                entity.set(KeyWordMqField.CREATION_DATE, current);
                if (CollUtil.isEmpty(tags)) {
                    entity.set(KeyWordMqField.CONSUMER_TAG, 1);
                } else  {
                    entity.set(KeyWordMqField.CONSUMER_TAG, tags.get(RandomUtil.randomInt(tags.size())));
                }
                list.add(entity);
            }
            db().tx(db ->  db().insert(list));
            return list.size();
        } catch (Exception e) {
            throw new DbMqException(e);
        }
    }

    @Override
    public List<ConsumerRegisterInfo> getRegisterInfo(List<String> topics) {
        try {
            return db().findAll(Entity.create(MySqlTableSql.TABLE_CONSUMER_REGISTER)
                    .set(KeyWordMqField.TOPIC, topics), ConsumerRegisterInfo.class);
        } catch (SQLException e) {
            throw new DbMqException(e);
        }
    }

    @Override
    public boolean tryGetConsumerRegister(ConsumerRegisterInfo registerInfo) {
        if (!Objects.equals(registerInfo.getStatus(), Const.Status.DOWN)) {
            return Boolean.FALSE;
        }
        String updateSql = "update " + MySqlTableSql.TABLE_CONSUMER_REGISTER +
                " set status = 'clear' " +
                " where id = ? and status = 'down'";
        try {
            final int result = db().execute(updateSql, registerInfo.getId());
            return result > 0 ;
        } catch (Exception e) {
            LOGGER.error(e.getMessage(), e);
            return false;
        }
    }

    @Override
    public void releaseConsumerRegister(ConsumerRegisterInfo registerInfo) {
        try {
            String updateSql = "update " + MySqlTableSql.TABLE_CONSUMER_REGISTER +
                    " set status = 'down' " +
                    " where id = ? and status = 'clear'";
            db().execute(updateSql, registerInfo.getId());
        } catch (Exception e) {
            LOGGER.error(e.getMessage(), e);
        }
    }

    @Override
    public void deleteConsumerRegister(ConsumerRegisterInfo registerInfo) {
        try {
            String updateSql = "delete from " + MySqlTableSql.TABLE_CONSUMER_REGISTER +
                    " where id = ? and status = 'clear'";
            db().execute(updateSql, registerInfo.getId());
        } catch (Exception e) {
            LOGGER.error(e.getMessage(), e);
        }
    }

    @Override
    public List<Message> getReAssignMessage(String topic, Integer tag) {
        try {
            return db().query("select * from `" + topic + "` where consumer_tag = ? and consume_flag = ? limit ?", Message.class, tag, 0, COUNT);
        } catch (SQLException e) {
            LOGGER.error(e.getMessage(), e);
        }
        return CollUtil.newArrayList();
    }

    @Override
    public void batchUpdateMessage(List<Message> messageList, List<Integer> tags, String topic) {
        // 在事务下批量更新，效率更高
        try {
            db().tx(db -> {
                for (Message message : messageList) {
                    Integer newTag = Util.getNewTag(tags);
                    db.execute("update `" + topic + "` set consumer_tag = ? where id = ?", newTag, message.getId());
                }
            });
        } catch (Exception e) {
            LOGGER.error(e.getMessage(), e);
        }
    }


}
